first commit

cv-model/main.py

@@ -0,0 +1,42 @@
+import torch
+from torchvision.models.detection import (
+    fasterrcnn_resnet50_fpn,
+    FasterRCNN_ResNet50_FPN_Weights
+)
+from PIL import Image
+import torchvision.transforms as T
+
+# Load pretrained model
+weights = FasterRCNN_ResNet50_FPN_Weights.DEFAULT
+model = fasterrcnn_resnet50_fpn(weights=weights)
+
+model.eval()
+
+# Load image
+image = Image.open("image.jpg").convert("RGB")
+
+# Convert image to tensor
+transform = T.ToTensor()
+image_tensor = transform(image)
+
+# Run detection
+with torch.no_grad():
+    predictions = model([image_tensor])
+
+prediction = predictions[0]
+
+# COCO class names
+categories = weights.meta["categories"]
+
+# Print detected objects
+for box, label, score in zip(
+    prediction["boxes"],
+    prediction["labels"],
+    prediction["scores"]
+):
+    if score > 0.5:
+        print(
+            f"{categories[label]}: "
+            f"{score:.2f}, "
+            f"box={box.tolist()}"
+        )