minor changes

model/pothole_model.py

@@ -1,44 +1,68 @@
-from keras.models import load_model  # TensorFlow is required for Keras to work
-from PIL import Image, ImageOps  # Install pillow instead of PIL
-import numpy as np
+# from ultralyticsplus import YOLO, render_result
+# from PIL import Image
+# 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_image_model(image):
-    np.set_printoptions(suppress=True)
+#     # 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 the model
-    model = load_model("keras_Model.h5", compile=False)
+#     # Convert PIL image to bytes
+#     # with io.BytesIO() as output:
+#     #     pil_image.save(output, format='JPEG')
+#     #     image_bytes = output.getvalue()
 
-    # Load the labels
-    class_names = open("labels.txt", "r").readlines()
+#     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
 
-    # Create the array of the right shape to feed into the keras model
-    # The 'length' or number of images you can put into the array is
-    # determined by the first position in the shape tuple, in this case 1
-    data = np.ndarray(shape=(1, 224, 224, 3), dtype=np.float32)
 
-    # Replace this with the path to your image
-    image = Image.open(image).convert("RGB")
+from PIL import Image
+from io import BytesIO
 
-    # resizing the image to be at least 224x224 and then cropping from the center
-    size = (224, 224)
-    image = ImageOps.fit(image, size, Image.Resampling.LANCZOS)
+# Load model directly
+from transformers import AutoImageProcessor, AutoModelForObjectDetection
 
-    # turn the image into a numpy array
-    image_array = np.asarray(image)
+processor = AutoImageProcessor.from_pretrained("savioratharv/pothole_detection")
+model = AutoModelForObjectDetection.from_pretrained("savioratharv/pothole_detection")
 
-    # Normalize the image
-    normalized_image_array = (image_array.astype(np.float32) / 127.5) - 1
+# 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")
 
-    # Load the image into the array
-    data[0] = normalized_image_array
+    # 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()
 
-    # Predicts the model
-    prediction = model.predict(data)
-    index = np.argmax(prediction)
-    class_name = class_names[index]
-    confidence_score = prediction[0][index]
+    return predicted_class
 
-    # Print prediction and confidence score
-    print("Class:", class_name[2:], end="")
-    print("Confidence Score:", confidence_score)

routes/route.py

@@ -1,7 +1,7 @@
 from fastapi import APIRouter, HTTPException,Depends,File, UploadFile
 from fastapi.responses import JSONResponse
 from config.database import admin_collection, user_collection,notification_collection,pothole_image_collection
-from model.pothole_model import load_image_model
+from model.pothole_model import predict_pothole
 from utils.auth import create_access_token, hash_password, verify_password, verify_token
 from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from schema.model import Admin, PoholeInfo, PotInfoById, PotholeFilters, PotholeModel, UpdatePotholeInfo, User, UserLogin, VerifyOtp
@@ -262,7 +262,7 @@ async def verify_pothole(potholeModel: PotholeModel, token: HTTPAuthorizationCre
         image_bytes = response.content
 
         # Pass image bytes to your model function
-        results = load_image_model(image_bytes)
+        results = predict_pothole(image_bytes)
 
         # if results == 1:
         #     return JSONResponse(content={"response": "Pothole"})