add endpoint for single url

main.py

@@ -16,6 +16,7 @@ from models import (
     FileImageDetectionResponse,
     UrlImageDetectionResponse,
     ImageUrlsRequest,
+    ImageUrlRequest,
 )
 
 app = FastAPI()
@@ -87,36 +88,9 @@ async def classify_image(file: UploadFile = File(None)):
         # inputs = image_processor(image, return_tensors="pt")
         inputs = model(image)
         logging.info("inputs %s", inputs)
-        # with torch.no_grad():
-        #     outpus = model(**inputs)
-        # logits = outpus.logits
-        # logging.info("logits %s", logits)
-        # probs = F.softmax(logits, dim=1)
-        # logging.info("probs %s", probs)
-        # predicted_label_id = probs.argmax(-1).item()
-        # logging.info("predicted_label_id %s", predicted_label_id)
-        # predicted_label = model.config.id2label[predicted_label_id]
-        # logging.info("model.config.id2label %s", model.config.id2label)
-        # confidence = probs.max().item()
-# outpus = model(**inputs)
-#             logits = outpus.logits
-#             probs = F.softmax(logits, dim=-1)
-#             predicted_label_id = probs.argmax(-1).item()
-#             predicted_label = model.config.id2label[predicted_label_id]
-#             confidence = probs.max().item()
-
-    # model predicts one of the 1000 ImageNet classes
-    #     predicted_label = logits.argmax(-1).item()
-    #     logging.info("predicted_label", predicted_label)
-    #     logging.info("model.config.id2label[predicted_label] %s", model.config.id2label[predicted_label])
-    # # print(model.config.id2label[predicted_label])
-    # Find the prediction with the highest confidence using the max() function
+ 
         predicted_label = max(inputs, key=lambda x: x["score"])
-    # logging.info("best_prediction %s", best_prediction)
-    # best_prediction2 = results[1]["label"]
-    # logging.info("best_prediction2 %s", best_prediction2)
-
-    # # Calculate the confidence score, rounded to the nearest tenth and as a percentage
+ 
         confidence = round(predicted_label["score"] * 100, 1)
 
     # # Prepare the custom response data
@@ -125,32 +99,13 @@ async def classify_image(file: UploadFile = File(None)):
             "prediction": predicted_label["label"],
             "confidence": str(confidence),
         }
-        # Use the model to classify the image
-        # results = model(image)
-
-        # Find the prediction with the highest confidence using the max() function
-        # best_prediction = max(results, key=lambda x: x["score"])
-
-        # Calculate the confidence score, rounded to the nearest tenth and as a percentage
-        # confidence_percentage = round(best_prediction["score"] * 100, 1)
-
-        # Prepare the custom response data
-        # detection_result = {
-        #     "is_nsfw": best_prediction["label"] == "nsfw",
-        #     "confidence_percentage": confidence_percentage,
-        # }
-
+    
         # Populate hash
         cache[image_hash] = response_data.copy()
 
         # Add url to the API response
         response_data["file_name"] = file.filename
 
-        # response_data.append(detection_result)
-
-        # Add file_name to the API response
-        # response_data["file_name"] = file.filename
-
         return FileImageDetectionResponse(**response_data)
 
     # except Exception as e:
@@ -161,6 +116,52 @@ async def classify_image(file: UploadFile = File(None)):
         ) from e
 
 
+@app.post("/v1/detect/url", response_model=UrlImageDetectionResponse)
+async def classify_images(request: ImageUrlRequest):
+ 
+    try:
+        image_url = request.url
+        logging.info("Downloading image from URL: %s", image_url)
+        image_data = await download_image(image_url)
+        image_hash = hash_data(image_data)
+
+        if image_hash in cache:
+            # Return cached entry
+            logging.info("Returning cached entry for %s", image_url)
+
+            cached_response = cache[image_hash]
+            response_data = {**cached_response, "url": image_url}
+
+            return UrlImageDetectionResponse(**response_data)
+
+        image = Image.open(io.BytesIO(image_data))
+        # inputs = image_processor(image, return_tensors="pt")
+        inputs = model(image)
+
+        predicted_label = max(inputs, key=lambda x: x["score"])
+        confidence = round(predicted_label["score"] * 100, 1)
+
+        response_data = {
+            "prediction": predicted_label["label"],
+            "confidence": str(confidence),
+        }
+  
+        # Populate hash
+        cache[image_hash] = response_data.copy()
+
+        # Add url to the API response
+        response_data["url"] = image_url
+
+        return UrlImageDetectionResponse(**response_data)
+
+    # except Exception as e:
+    except PipelineException as e:
+        logging.error("Error processing image from %s: %s", image_url, str(e))
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error processing image from {image_url}: {str(e)}",
+        ) from e
+
 @app.post("/v1/detect/urls", response_model=list[UrlImageDetectionResponse])
 async def classify_images(request: ImageUrlsRequest):
     """Function analyzing images from URLs."""

models.py

@@ -3,6 +3,16 @@
 from pydantic import BaseModel
 
 
+class ImageUrlRequest(BaseModel):
+    """
+    Model representing the request body for the /v1/detect/urls endpoint.
+
+    Attributes:
+        urls (list[str]): List of image URLs to be processed.
+    """
+
+    url: str
+
 class ImageUrlsRequest(BaseModel):
     """
     Model representing the request body for the /v1/detect/urls endpoint.