minor changes

model/pothole_model.py

@@ -1,68 +1,69 @@
-# from ultralyticsplus import YOLO, render_result
-# from PIL import Image
-# import numpy as np
+from ultralyticsplus import YOLO, render_result
+from PIL import Image
+from io import BytesIO
+import numpy as np
 
-# def load_model(image):
-#     # image_bytes = image.content
-#     model = YOLO('keremberke/yolov8n-pothole-segmentation')
-#     model.overrides['conf'] = 0.25
-#     model.overrides['iou'] = 0.45
-#     model.overrides['agnostic_nms'] = False
-#     model.overrides['max_det'] = 1000
+def load_model(image):
+    # image_bytes = image.content
+    model = YOLO('keremberke/yolov8n-pothole-segmentation')
+    model.overrides['conf'] = 0.25
+    model.overrides['iou'] = 0.45
+    model.overrides['agnostic_nms'] = False
+    model.overrides['max_det'] = 1000
 
-#     # Load image using PIL
-#     image = Image.open((image))
-#     image_array = np.array(image)
-#     # pil_image = pil_image.convert("RGB")  # Ensure image is in RGB format
+    # Load image using PIL
+    image = Image.open(BytesIO(image))
+    image_array = np.array(image)
+    # pil_image = pil_image.convert("RGB")  # Ensure image is in RGB format
 
-#     # Convert PIL image to bytes
-#     # with io.BytesIO() as output:
-#     #     pil_image.save(output, format='JPEG')
-#     #     image_bytes = output.getvalue()
+    # Convert PIL image to bytes
+    # with io.BytesIO() as output:
+    #     pil_image.save(output, format='JPEG')
+    #     image_bytes = output.getvalue()
 
-#     results = model.predict(image_array)
-#     for result in results:
-#         boxes = result.boxes.xyxy
-#         conf = result.boxes.conf
-#         cls = result.boxes.cls
-#         obj_info = []
-#         for i, bbox in enumerate(boxes):
-#             label = result.names[int(cls[i])]
-#             obj_info.append({
-#                 "Object": i+1,
-#                 "Label": label,
-#                 "Confidence": conf[i],
-#                 "Bounding Box": bbox
-#             })
-#     render = render_result(model=model, image=image, result=results[0])
-#     if label:
-#         print(label)
-#     render.show()
-#     return label
+    results = model.predict(image_array)
+    for result in results:
+        boxes = result.boxes.xyxy
+        conf = result.boxes.conf
+        cls = result.boxes.cls
+        obj_info = []
+        for i, bbox in enumerate(boxes):
+            label = result.names[int(cls[i])]
+            obj_info.append({
+                "Object": i+1,
+                "Label": label,
+                "Confidence": conf[i],
+                "Bounding Box": bbox
+            })
+    render = render_result(model=model, image=image, result=results[0])
+    if label:
+        print(label)
+    render.show()
+    return label
 
 
-from PIL import Image
-from io import BytesIO
+# from PIL import Image
+# from io import BytesIO
 
-# Load model directly
-from transformers import AutoImageProcessor, AutoModelForObjectDetection
+# # Load model directly
+# from transformers import AutoImageProcessor, AutoModelForObjectDetection
 
-processor = AutoImageProcessor.from_pretrained("savioratharv/pothole_detection")
-model = AutoModelForObjectDetection.from_pretrained("savioratharv/pothole_detection")
+# processor = AutoImageProcessor.from_pretrained("savioratharv/pothole_detection")
+# model = AutoModelForObjectDetection.from_pretrained("savioratharv/pothole_detection")
 
-# Function to predict if an image contains a pothole
-def predict_pothole(image_url):
-    image = Image.open(BytesIO(image_url))
-    inputs = processor(images=image, return_tensors="pt")
+# # Function to predict if an image contains a pothole
+# def predict_pothole(image_url):
+#     image = Image.open(BytesIO(image_url))
+#     inputs = processor(images=image, return_tensors="pt")
 
-    # Perform inference
-    outputs = model(**inputs)
-    logits = outputs.logits
-    probabilities = logits.softmax(dim=1)
+#     # Perform inference
+#     outputs = model(**inputs)
+#     logits = outputs.logits
+#     probabilities = logits.softmax(dim=1)
     
-    # Get predicted class (0: No pothole, 1: Pothole)
-    predicted_class = probabilities.argmax().item()
-    confidence = probabilities[0, predicted_class].item()
+#     # Get predicted class (0: No pothole, 1: Pothole)
+#     predicted_class = probabilities.argmax().item()
+#     confidence = probabilities[0, predicted_class].item()
 
-    return predicted_class
+#     return predicted_class
 

requirements.txt

@@ -8,5 +8,4 @@ python-multipart
 certifi
 firebase_admin
 pillow
-transformers
-pytorch-lightning
+ultralyticsplus