Uploading FoodExtract-Vision demo app.py from YouTube tutorial video

app.py

@@ -1,78 +1,185 @@
 import torch
 import gradio as gr
+
 import spaces
 from transformers import pipeline
 
 BASE_MODEL_ID = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct"
+# FINE_TUNED_MODEL_ID = "mrdbourke/FoodExtract-Vision-SmolVLM2-500M-fine-tune-v1"
 FINE_TUNED_MODEL_ID = "ninjals/FoodExtract-Vision-SmolVLM2-500M-fine-tune-v1-VIDEO"
 OUTPUT_TOKENS = 256
 
-DTYPE = torch.float16  # ✅ safest default
-
-original_pipeline = None
-ft_pipe = None
-
-def _load_pipes():
-    global original_pipeline, ft_pipe
-    if original_pipeline is None:
-        print("[INFO] Loading Original Model")
-        original_pipeline = pipeline(
-            "image-text-to-text",
-            model=BASE_MODEL_ID,
-            device_map="auto",
-            dtype=DTYPE,
-        )
-    if ft_pipe is None:
-        print("[INFO] Loading Fine-tuned Model")
-        ft_pipe = pipeline(
-            "image-text-to-text",
-            model=FINE_TUNED_MODEL_ID,
-            device_map="auto",
-            dtype=DTYPE,  # ✅ keep fp16 (don’t use bf16)
-        )
+# Load original base model (no fine-tuning)
+print(f"[INFO] Loading Original Model")
+original_pipeline = pipeline(
+    "image-text-to-text",
+    model=BASE_MODEL_ID,
+    dtype=torch.float16,
+    device_map="auto"
+)
+
+# Load fine-tuned model
+print(f"[INFO] Loading Fine-tuned Model")
+ft_pipe = pipeline(
+    "image-text-to-text",
+    model=FINE_TUNED_MODEL_ID,
+    dtype=torch.bfloat16,
+    device_map="auto"
+)
 
 def create_message(input_image):
-    prompt = (
-        "Classify the given input image into food or not and if edible food or drink items are present, "
-        "extract those to a list. If no food/drink items are visible, return empty lists.\n\n"
-        'Only return valid JSON in the following form:\n\n'
-        '{\n'
-        '  "is_food": 0,\n'
-        '  "image_title": "",\n'
-        '  "food_items": [],\n'
-        '  "drink_items": []\n'
-        '}\n'
-    )
-
-    return [{
-        "role": "user",
-        "content": [
-            {"type": "image", "image": input_image},
-            {"type": "text", "text": prompt},
-        ],
-    }]
+    return [{'role': 'user',
+ 'content': [{'type': 'image',
+   'image': input_image},
+  {'type': 'text',
+   'text': "Classify the given input image into food or not and if edible food or drink items are present, extract those to a list. If no food/drink items are visible, return empty lists.\n\nOnly return valid JSON in the following form:\n\n```json\n{\n  'is_food': 0, # int - 0 or 1 based on whether food/drinks are present (0 = no foods visible, 1 = foods visible)\n  'image_title': '', # str - short food-related title for what foods/drinks are visible in the image, leave blank if no foods present\n  'food_items': [], # list[str] - list of visible edible food item nouns\n  'drink_items': [] # list[str] - list of visible edible drink item nouns\n}\n```\n"}]}]
 
 @spaces.GPU
 def extract_foods_from_image(input_image):
-    _load_pipes()
+    input_image = input_image.resize(size=(512, 512))
+    input_message = create_message(input_image=input_image)
 
-    input_image = input_image.resize((512, 512))
-    input_message = create_message(input_image)
+    # Get outputs from base model (not fine-tuned)
+    original_pipeline_output = original_pipeline(text=[input_message],
+                                                 max_new_tokens=OUTPUT_TOKENS)
 
-    base_out = original_pipeline(text=[input_message], max_new_tokens=OUTPUT_TOKENS, do_sample=False)
-    ft_out   = ft_pipe(text=[input_message], max_new_tokens=OUTPUT_TOKENS, do_sample=False)
+    outputs_pretrained = original_pipeline_output[0][0]["generated_text"][-1]["content"]
+
+    # Get outputs from fine-tuned model (fine-tuned on food images)
+    ft_pipe_output = ft_pipe(text=[input_message],
+                             max_new_tokens=OUTPUT_TOKENS)
+    outputs_fine_tuned = ft_pipe_output[0][0]["generated_text"][-1]["content"]
 
-    outputs_pretrained = base_out[0][0]["generated_text"][-1]["content"]
-    outputs_fine_tuned = ft_out[0][0]["generated_text"][-1]["content"]
     return outputs_pretrained, outputs_fine_tuned
 
+demo_title = "🥑➡️📝 FoodExtract-Vision with a fine-tuned SmolVLM2-500M"
+demo_description = """* **Base model:** https://huggingface.co/HuggingFaceTB/SmolVLM-500M-Instruct
+* **Fine-tuning dataset:** https://huggingface.co/datasets/mrdbourke/FoodExtract-1k-Vision (1k food images and 500 not food images)
+* **Fine-tuned model:** https://huggingface.co/ninjals/FoodExtract-Vision-SmolVLM2-500M-fine-tune-v1-VIDEO
+
+## Overview
+
+Extract food and drink items in a structured way from images.
+
+The original model outputs fail to capture the desired structure. But the fine-tuned model sticks to the output structure quite well.
+
+However, the fine-tuned model could definitely be improved with respects to its ability to extract the right food/drink items.
+
+Both models use the input prompt:
+
+````
+Classify the given input image into food or not and if edible food or drink items are present, extract those to a list. If no food/drink items are visible, return empty lists.
+
+Only return valid JSON in the following form:
+
+```json
+{
+  'is_food': 0, # int - 0 or 1 based on whether food/drinks are present (0 = no foods visible, 1 = foods visible)
+  'image_title': '', # str - short food-related title for what foods/drinks are visible in the image, leave blank if no foods present
+  'food_items': [], # list[str] - list of visible edible food item nouns
+  'drink_items': [] # list[str] - list of visible edible drink item nouns
+}
+```
+````
+
+Except one model has been fine-tuned on the structured data whereas the other hasn't.
+
+Notable next steps would be:
+* **Remove the input prompt:** Just train the model to go straight from image -> text (no text prompt on input), this would save on inference tokens.
+* **Fine-tune on more real-world data:** Right now the model is only trained on 1k food images (from Food101) and 500 not food (random internet images), training on real world data would likely significantly improve performance.
+* **Fix the repetitive generation:** The model can sometimes get stuck in a repetitive generation pattern, e.g. "onions", "onions", "onions", etc. We could look into patterns to help reduce this.
+"""
+
 demo = gr.Interface(
     fn=extract_foods_from_image,
     inputs=gr.Image(type="pil"),
-    outputs=[gr.Textbox(lines=4), gr.Textbox(lines=4)],
-    examples=[["examples/camera.jpeg"], ["examples/Tandoori-Chicken.jpg"], ["examples/fries.jpeg"]],
+    title=demo_title,
+    description=demo_description,
+    outputs=[gr.Textbox(lines=4, label="Original Model (not fine-tuned)"),
+             gr.Textbox(lines=4, label="Fine-tuned Model")],
+    examples=[["examples/camera.jpeg"],
+              ["examples/Tandoori-Chicken.jpg"],
+              ["examples/fries.jpeg"]],
 )
 
 if __name__ == "__main__":
-    demo.launch(show_error=True)
+    demo.launch(share=False, show_error=True)
+
+# %%writefile demos/FoodExtract-Vision-v1/app.py
+# import torch
+# import gradio as gr
+# import spaces
+# from transformers import pipeline
+
+# BASE_MODEL_ID = "HuggingFaceTB/SmolVLM2-500M-Video-Instruct"
+# FINE_TUNED_MODEL_ID = "ninjals/FoodExtract-Vision-SmolVLM2-500M-fine-tune-v1-VIDEO"
+# OUTPUT_TOKENS = 256
+
+# DTYPE = torch.float16  # ✅ safest default
+
+# original_pipeline = None
+# ft_pipe = None
+
+# def _load_pipes():
+#     global original_pipeline, ft_pipe
+#     if original_pipeline is None:
+#         print("[INFO] Loading Original Model")
+#         original_pipeline = pipeline(
+#             "image-text-to-text",
+#             model=BASE_MODEL_ID,
+#             device_map="auto",
+#             dtype=DTYPE,
+#         )
+#     if ft_pipe is None:
+#         print("[INFO] Loading Fine-tuned Model")
+#         ft_pipe = pipeline(
+#             "image-text-to-text",
+#             model=FINE_TUNED_MODEL_ID,
+#             device_map="auto",
+#             dtype=DTYPE,  # ✅ keep fp16 (don’t use bf16)
+#         )
+
+# def create_message(input_image):
+#     prompt = (
+#         "Classify the given input image into food or not and if edible food or drink items are present, "
+#         "extract those to a list. If no food/drink items are visible, return empty lists.\n\n"
+#         'Only return valid JSON in the following form:\n\n'
+#         '{\n'
+#         '  "is_food": 0,\n'
+#         '  "image_title": "",\n'
+#         '  "food_items": [],\n'
+#         '  "drink_items": []\n'
+#         '}\n'
+#     )
+
+#     return [{
+#         "role": "user",
+#         "content": [
+#             {"type": "image", "image": input_image},
+#             {"type": "text", "text": prompt},
+#         ],
+#     }]
+
+# @spaces.GPU
+# def extract_foods_from_image(input_image):
+#     _load_pipes()
+
+#     input_image = input_image.resize((512, 512))
+#     input_message = create_message(input_image)
+
+#     base_out = original_pipeline(text=[input_message], max_new_tokens=OUTPUT_TOKENS, do_sample=False)
+#     ft_out   = ft_pipe(text=[input_message], max_new_tokens=OUTPUT_TOKENS, do_sample=False)
+
+#     outputs_pretrained = base_out[0][0]["generated_text"][-1]["content"]
+#     outputs_fine_tuned = ft_out[0][0]["generated_text"][-1]["content"]
+#     return outputs_pretrained, outputs_fine_tuned
+
+# demo = gr.Interface(
+#     fn=extract_foods_from_image,
+#     inputs=gr.Image(type="pil"),
+#     outputs=[gr.Textbox(lines=4), gr.Textbox(lines=4)],
+#     examples=[["examples/camera.jpeg"], ["examples/Tandoori-Chicken.jpg"], ["examples/fries.jpeg"]],
+# )
 
+# if __name__ == "__main__":
+#     demo.launch(show_error=True)