Update app.py

app.py

@@ -1,55 +1,30 @@
-import os, json
-import numpy as np
-import tensorflow as tf
-from PIL import Image
 import gradio as gr
-from huggingface_hub import snapshot_download
-
-# 1️⃣ Download the entire SavedModel repo
-REPO_ID   = "Adriana213/vgg16-fruit-classifier"
-local_dir = snapshot_download(repo_id=REPO_ID)
-
-# 2️⃣ Load labels
-with open(os.path.join(local_dir, "class_labels.json"), "r") as f:
-    id2label = json.load(f)
-
-# 3️⃣ Load as a TF SavedModel and grab the default signature
-loaded = tf.saved_model.load(local_dir)
-infer   = loaded.signatures["serving_default"]
-
-# 4️⃣ Figure out the input key name
-input_key = list(infer.structured_input_signature[1].keys())[0]
-
-# 5️⃣ Preprocess helper (same as before)
-def preprocess(img: Image.Image) -> np.ndarray:
-    img = img.resize((100,100))
-    arr = np.array(img.convert("RGB")).astype("float32")
-    arr[...,0] -= 123.68
-    arr[...,1] -= 116.779
-    arr[...,2] -= 103.939
-    return np.expand_dims(arr, 0)
-
-# 6️⃣ Inference fn
-def classify_fruit(img: Image.Image):
-    x = preprocess(img)
-    # call the signature with the correct kwarg
-    outputs = infer(**{input_key: tf.constant(x)})
-    # grab the first tensor in the outputs dict
-    logits = list(outputs.values())[0].numpy()
-    probs  = tf.nn.softmax(logits, axis=1).numpy()
-    idx    = int(np.argmax(probs, axis=1)[0])
-    label  = id2label[str(idx)]
-    score  = float(probs[0][idx])
-    return {label: round(score, 4)}
+import torch
+from transformers import AutoModelForImageClassification, AutoImageProcessor
+from PIL import Image
 
-# 7️⃣ Gradio interface
-demo = gr.Interface(
-    fn=classify_fruit,
-    inputs=gr.Image(type="pil", label="Upload Fruit Image"),
-    outputs=gr.Label(num_top_classes=5, label="Top Predictions"),
-    title="🍉 Fruit Classifier (131 types)",
-    description="Classify a fruit image into one of 131 categories."
+# Load model and processor
+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_image(image):
+    # Preprocess the image
+    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
+
+# Create Gradio interface
+interface = gr.Interface(
+    fn=classify_image,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Text(label="Detected Label"),
+    title="Fruit & Vegetable Detector",
+    description="Upload an image of a fruit or vegetable and get the predicted label."
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    interface.launch()