PyCharm Code Inspect fixes for Solutions and Examples (#18393)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Muhammad Rizwan Munawar 2024-12-25 23:36:35 +05:00 committed by GitHub
parent 3e65fc2421
commit b1af683d7a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 92 additions and 62 deletions

View file

@ -132,17 +132,19 @@ def run(
model.to("cuda") if device == "0" else model.to("cpu")
# Extract classes names
names = model.model.names
names = model.names
# Video setup
videocapture = cv2.VideoCapture(source)
frame_width, frame_height = int(videocapture.get(3)), int(videocapture.get(4))
fps, fourcc = int(videocapture.get(5)), cv2.VideoWriter_fourcc(*"mp4v")
frame_width = int(videocapture.get(3))
frame_height = int(videocapture.get(4))
fps = int(videocapture.get(5))
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
# Output setup
save_dir = increment_path(Path("ultralytics_rc_output") / "exp", exist_ok)
save_dir.mkdir(parents=True, exist_ok=True)
video_writer = cv2.VideoWriter(str(save_dir / f"{Path(source).stem}.mp4"), fourcc, fps, (frame_width, frame_height))
video_writer = cv2.VideoWriter(str(save_dir / f"{Path(source).stem}.avi"), fourcc, fps, (frame_width, frame_height))
# Iterate over video frames
while videocapture.isOpened():
@ -241,9 +243,9 @@ def parse_opt():
return parser.parse_args()
def main(opt):
def main(options):
"""Main function."""
run(**vars(opt))
run(**vars(options))
if __name__ == "__main__":

View file

@ -24,11 +24,16 @@ class SAHIInference:
yolo11_model_path = f"models/{weights}"
download_yolo11n_model(yolo11_model_path)
self.detection_model = AutoDetectionModel.from_pretrained(
model_type="ultralytics", model_path=yolo11_model_path, confidence_threshold=0.3, device="cpu"
model_type="ultralytics", model_path=yolo11_model_path, device="cpu"
)
def inference(
self, weights="yolo11n.pt", source="test.mp4", view_img=False, save_img=False, exist_ok=False, track=False
self,
weights="yolo11n.pt",
source="test.mp4",
view_img=False,
save_img=False,
exist_ok=False,
):
"""
Run object detection on a video using YOLO11 and SAHI.
@ -39,7 +44,6 @@ class SAHIInference:
view_img (bool): Show results.
save_img (bool): Save results.
exist_ok (bool): Overwrite existing files.
track (bool): Enable object tracking with SAHI
"""
# Video setup
cap = cv2.VideoCapture(source)
@ -50,8 +54,8 @@ class SAHIInference:
save_dir = increment_path(Path("ultralytics_results_with_sahi") / "exp", exist_ok)
save_dir.mkdir(parents=True, exist_ok=True)
video_writer = cv2.VideoWriter(
str(save_dir / f"{Path(source).stem}.mp4"),
cv2.VideoWriter_fourcc(*"mp4v"),
str(save_dir / f"{Path(source).stem}.avi"),
cv2.VideoWriter_fourcc(*"MJPG"),
int(cap.get(5)),
(frame_width, frame_height),
)
@ -68,8 +72,6 @@ class SAHIInference:
self.detection_model,
slice_height=512,
slice_width=512,
overlap_height_ratio=0.2,
overlap_width_ratio=0.2,
)
detection_data = [
(det.category.name, det.category.id, (det.bbox.minx, det.bbox.miny, det.bbox.maxx, det.bbox.maxy))