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:
parent
3e65fc2421
commit
b1af683d7a
13 changed files with 92 additions and 62 deletions
|
|
@ -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__":
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue