Classifying common fast-food items

app.py

@@ -23,57 +23,8 @@ def transform_img(img: Image.Image) -> torch.tensor:
 
 # Returns string with class and probability
 def classify_img(img: Image.Image) -> dict:
-    class_names = ['AIR COMPRESSOR',
- 'ALTERNATOR',
- 'BATTERY',
- 'BRAKE CALIPER',
- 'BRAKE PAD',
- 'BRAKE ROTOR',
- 'CAMSHAFT',
- 'CARBERATOR',
- 'CLUTCH PLATE',
- 'COIL SPRING',
- 'CRANKSHAFT',
- 'CYLINDER HEAD',
- 'DISTRIBUTOR',
- 'ENGINE BLOCK',
- 'ENGINE VALVE',
- 'FUEL INJECTOR',
- 'FUSE BOX',
- 'GAS CAP',
- 'HEADLIGHTS',
- 'IDLER ARM',
- 'IGNITION COIL',
- 'INSTRUMENT CLUSTER',
- 'LEAF SPRING',
- 'LOWER CONTROL ARM',
- 'MUFFLER',
- 'OIL FILTER',
- 'OIL PAN',
- 'OIL PRESSURE SENSOR',
- 'OVERFLOW TANK',
- 'OXYGEN SENSOR',
- 'PISTON',
- 'PRESSURE PLATE',
- 'RADIATOR',
- 'RADIATOR FAN',
- 'RADIATOR HOSE',
- 'RADIO',
- 'RIM',
- 'SHIFT KNOB',
- 'SIDE MIRROR',
- 'SPARK PLUG',
- 'SPOILER',
- 'STARTER',
- 'TAILLIGHTS',
- 'THERMOSTAT',
- 'TORQUE CONVERTER',
- 'TRANSMISSION',
- 'VACUUM BRAKE BOOSTER',
- 'VALVE LIFTER',
- 'WATER PUMP',
- 'WINDOW REGULATOR']
-    model = torch.jit.load("car_part_traced_classifier_resnet50.ptl")
+    class_names = ['Baked Potato','Burger','Crispy Chicken','Donut','Fries','Hot Dog','Pizza','Sandwich','Taco','Taquito']
+    model = torch.jit.load("fast_food_traced_classifier_resnet50.ptl")
     # Applying transformation to the image
     model_img = transform_img(img)
     model_img = model_img.view(1,3,224,224)
@@ -91,10 +42,10 @@ def classify_img(img: Image.Image) -> dict:
     # Dictionary I will display
     model_output = {}
     for i in range(top3_prob.size(1)):
-        car_part_name = class_names[top3_catid[0][i].item()]
+        fast_food_name = class_names[top3_catid[0][i].item()]
         probability = round(top3_prob[0][i].item() * 100, 2)
         
-        model_output[f"top{i+1}"] = {"name": car_part_name, "probability": probability}
+        model_output[f"top{i+1}"] = {"name": fast_food_name, "probability": probability}
     
     return model_output