Uploaded new version

Dockerfile

@@ -17,12 +17,6 @@ USER user
 # Set home to the user's home directory
 ENV HOME=/home/user \
     PATH=/home/user/.local/bin:$PATH
-
-
-# Define environment variables for GPU, RAM limit, and shared memory
-ENV GPU_ID="0" \
-    RAM_LIMIT="4g" \
-    SHM_SIZE="2g"
  
 # Set the working directory to the user's home directory
 WORKDIR $HOME/app
@@ -31,4 +25,4 @@ WORKDIR $HOME/app
 COPY --chown=user . $HOME/app
  
 # Start the FastAPI app on port 7860, the default port expected by Spaces
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["uvicorn", "app-local:app", "--host", "0.0.0.0", "--port", "7860"]

app-hf.py

@@ -0,0 +1,53 @@
+from fastapi import FastAPI, HTTPException, Response
+from fastapi.responses import HTMLResponse
+from transformers import pipeline, YolosForObjectDetection, YolosImageProcessor
+from PIL import Image, ImageDraw
+import torch
+import requests
+import io
+import base64
+
+
+app = FastAPI() # Create a new FastAPI app instance
+
+# Initialize the Yolos model and image processor
+yolos_model = YolosForObjectDetection.from_pretrained('hustvl/yolos-tiny')
+yolos_image_processor = YolosImageProcessor.from_pretrained("hustvl/yolos-tiny")
+
+# Define a route for the root "/"
+@app.get("/")
+def read_root():
+    return {"message": "Welcome to the YOLOS Object Detection API!"}
+
+# Define a route for detecting objects from an image URL
+@app.get("/", response_class=HTMLResponse)
+def detect_objects(url: str):
+    try:
+        # Download the image from the specified URL
+        image = Image.open(requests.get(url, stream=True).raw)
+
+        # Preprocess the image using the Yolos image processor
+        inputs = yolos_image_processor(images=image, return_tensors="pt")
+
+        # Run the Yolos model on the preprocessed image
+        outputs = yolos_model(**inputs)
+
+        # Post-process the object detection results
+        target_sizes = torch.tensor([image.size[::-1]])
+        results = yolos_image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[0]
+
+        # Draw bounding boxes on the image
+        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+            image_draw = ImageDraw.Draw(image)
+            image_draw.rectangle(box.tolist(), outline="red", width=2)
+            image_draw.text((box[0], box[1]), f"{yolos_model.config.id2label[label.item()]}: {round(score.item(), 3)}", fill="red")
+
+        # Save the modified image to a byte stream
+        image_byte_array = io.BytesIO()
+        image.save(image_byte_array, format="PNG")
+
+        # Return the image as a Response with content type "image/png"
+        return Response(content=image_byte_array.getvalue(), media_type="image/png")
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")

app-local.py

@@ -0,0 +1,61 @@
+from fastapi import FastAPI, File, UploadFile, HTTPException, Form
+from fastapi.responses import HTMLResponse
+from transformers import YolosForObjectDetection, YolosImageProcessor
+from PIL import Image, ImageDraw
+import torch
+import io
+import base64
+from starlette.requests import Request
+from fastapi.templating import Jinja2Templates
+import httpx
+
+app = FastAPI() # Create a FastAPI instance
+templates = Jinja2Templates(directory="templates") # Create a Jinja2Templates instance for handling HTML templates
+
+# Initialize YOLOS model and image processor
+yolos_model = YolosForObjectDetection.from_pretrained('hustvl/yolos-tiny')
+yolos_image_processor = YolosImageProcessor.from_pretrained("hustvl/yolos-tiny")
+
+# Define a route for the main HTML page
+@app.get("/", response_class=HTMLResponse)
+async def main(request: Request):
+    return templates.TemplateResponse("index.html", {"request": request})
+
+# Define a route for handling object detection from a submitted form
+@app.post("/", response_class=HTMLResponse)
+async def post_detect_objects(request: Request, url: str = Form(...)):
+    try:
+        # Download the image from the specified URL
+        async with httpx.AsyncClient() as client:
+            response = await client.get(url)
+            response.raise_for_status()  # Raise an exception if there is an error in the request
+            content = response.content
+
+        image = Image.open(io.BytesIO(content))
+
+
+        # Preprocess the image using the YOLOS image processor
+        inputs = yolos_image_processor(images=image, return_tensors="pt")
+
+        # Run the YOLOS model on the preprocessed image
+        outputs = yolos_model(**inputs)
+
+        # Post-process the object detection results
+        target_sizes = torch.tensor([image.size[::-1]])
+        results = yolos_image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[0]
+
+        # Draw bounding boxes on the image
+        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+            image_draw = ImageDraw.Draw(image)
+            image_draw.rectangle(box.tolist(), outline="red", width=2)
+            image_draw.text((box[0], box[1]), f"{yolos_model.config.id2label[label.item()]}: {round(score.item(), 3)}", fill="red")
+
+        # Save the modified image to a byte sequence
+        image_byte_array = io.BytesIO()
+        image.save(image_byte_array, format="PNG")
+        
+        # Return the image as a response with content type "image/png"
+        return templates.TemplateResponse("result.html", {"request": request, "image": base64.b64encode(image_byte_array.getvalue()).decode()})
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")

requirements.txt

@@ -4,4 +4,8 @@ uvicorn[standard]==0.17.*
 sentencepiece==0.1.*
 torch==1.11.*
 transformers==4.*
-Pillow==8.3.2
+Pillow==8.2.*
+starlette==0.17.*
+httpx==0.21.*
+jinja2==3.0.*
+python-multipart==0.0.*

templates/index.html

@@ -0,0 +1,15 @@
+<!DOCTYPE html>
+<html>
+<head>
+    <title>YOLOS Object Detection</title>
+</head>
+<body>
+    <h1>YOLOS Object Detection</h1>
+    <form action="/" method="post">
+        <label for="url">Image URL:</label>
+        <input type="text" id="url" name="url" required>
+        <button type="submit">Detect Objects</button>
+    </form>
+</body>
+</html>
+

templates/result.html

@@ -0,0 +1,10 @@
+<!DOCTYPE html>
+<html>
+<head>
+    <title>YOLOS Object Detection Result</title>
+</head>
+<body>
+    <h1>YOLOS Object Detection Result</h1>
+    <img src="data:image/png;base64,{{ image }}" alt="Detected Objects">
+</body>
+</html>