Add application file

Dockerfile

@@ -0,0 +1,11 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+COPY requirements.txt .
+RUN pip install -r requirements.txt
+
+COPY ./app /app/app
+
+EXPOSE 7860
+
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

app/food_model.py

@@ -0,0 +1,18 @@
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+from PIL import Image
+import torch
+
+# Load model and feature extractor once
+model_name = "DrishtiSharma/finetuned-ViT-Indian-Food-Classification-v3"
+model = ViTForImageClassification.from_pretrained(model_name)
+extractor = ViTFeatureExtractor.from_pretrained(model_name)
+
+def classify_indian_food(image: Image.Image):
+    # Preprocess and predict
+    inputs = extractor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits
+        predicted_idx = logits.argmax(-1).item()
+        label = model.config.id2label[predicted_idx]
+    return label

app/fruit_model.py

@@ -0,0 +1,17 @@
+from transformers import AutoModelForImageClassification, AutoImageProcessor
+from PIL import Image
+import torch
+
+# Load model and processor once at startup
+model = AutoModelForImageClassification.from_pretrained("jazzmacedo/fruits-and-vegetables-detector-36")
+processor = AutoImageProcessor.from_pretrained("jazzmacedo/fruits-and-vegetables-detector-36")
+labels = list(model.config.id2label.values())
+
+def classify_food(image: Image.Image):
+    # Preprocess the image using the Hugging Face processor
+    inputs = processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    predicted_idx = torch.argmax(outputs.logits, dim=1).item()
+    predicted_label = labels[predicted_idx]
+    return predicted_label

app/main.py

@@ -0,0 +1,23 @@
+from fastapi import FastAPI, UploadFile, File
+from PIL import Image
+import io
+from .food_model import classify_food
+from .fruit_model import classify_fruit
+
+app = FastAPI()
+
+@app.get("/")
+def root():
+    return {"message": "Welcome to Fruit & Food Classifier API"}
+
+@app.post("/predict/food")
+async def predict_food(file: UploadFile = File(...)):
+    image = Image.open(io.BytesIO(await file.read())).convert("RGB")
+    result = classify_food(image)
+    return {"prediction": result}
+
+@app.post("/predict/fruit")
+async def predict_fruit(file: UploadFile = File(...)):
+    image = Image.open(io.BytesIO(await file.read())).convert("RGB")
+    result = classify_fruit(image)
+    return {"prediction": result}

requirements.txt

@@ -0,0 +1,5 @@
+fastapi
+uvicorn
+torch
+transformers
+Pillow