updated fixes

app.py

@@ -26,57 +26,33 @@ def import_model(model_name):
     
     return learn
 #########################################################################
-geo = {"Circles": "260101-260131", "Ovals": "260301-260328", "Triangles": "260501-260528", "Diamonds": "260701-260728", "Squares": "260901-260928", "Rectangles": "261101-261128", "Quadrilaterals": "261301-261328", "Polygons": "261501-261528", "Lines, bands, bars": "261701-261725", "Geometric solids": "261901-261925"}
-scen = {"Mountains, rocks, caves": "060101-060104", "Scenery with water, rivers or streams": "060301-060325", "Other scenery": "060501-060502, 060901-060925", "Urban scenery or village scenes":  "060701-060703", "Stars, comets": "010101-010114",  "Planets, asteroids, meteors, moons": "010301-010304, 010901-010925, 011101-011125" ,"Sun": "010501-010525",  "Globes": "010701-010725", "Natural Phenomena": "011501-011525", "Maps or outlines of continents, countries and other geographical areas": "011701-011725"}
-char = {"Men": "020101-020139", "Women": "020301-020328", "Children": "020501-020527", "Groups of humans": "020701-020726", "Humans depicted engaging in activities": "020901-020919", "Parts of the human body, skeletons, skulls": "021101-021125", "Winged personages, fairies, supernatural beings": "040101-040125", "Beings partly human and partly animal": "040301-040325", "Mythological or legendary animals": "040501-040525", "Plants, objects or geometric figures representing a person or an animal": "040701-040707", "Masks": "40901"}
-misc = {"Letters or numerals including punctuation, forming figurative elements": "270101-270301", "Figurative elements forming representations of letters or numerals, including punctuation": "270301-270501", "Illegible signatures": "270501", "Inscriptions": "280101-280201", "Other forms of communication": "280201", "Miscellaneous": "290107", "Red or pink": "29-01", "Brown": "29-02", "Blue": "29-03", "Gray or silver": "29-04", "Violet or purple": "29-05", "Green": "29-06", "Orange": "29-07", "Yellow or gold": "29-08", "White": "29-09", "Clear or translucent": "29-10", "Black": "29-11"}
-build = {"Dwellings, cages or kennels": "70101-70301", "Buildings": "70301-70501", "Interior/Exterior parts of dwellings or buildings": "70501-70901", "Monuments, stadiums, fountains": "70901-71101", "Structural works": "71101-71301", "Billboards, signs": "71301-71501", "Building Materials": "71501-71525", "Furniture": "120101-120301", "Electrical equipment": "130101-130301", "Machines, appliances": "130301-130325", "Computer devices and office and business machines": "150101-150701", "Wheels, bearings": "150701-150901"}
-anim = {"Cats, dogs, wolves, foxes, bears, lions, tigers": "30101-30301", "Elephants, hippopotami, rhinoceri, giraffes, alpacas, camels, llamas": "30301-30501", "Horses, donkeys, zebras": "30501-30701", "Bovines, deer, antelopes, goats, sheep, pigs, cows, bulls, buffalo, moose": "30701-30901", "Small mammals other than cats and dogs, rodents, kangaroos and wallabies": "30901-31101", "Primates, (monkeys, apes, etc.)": "31101-31301", "Parts of the body, animal skeletons, animal skulls": "31301-31501", "Birds, bats": "31501-31701", "Parts of birds, eggs and nests": "31701-31901", "Fish, whales, seals, sea lions": "31901-32101", "Reptiles, snails, frogs": "32101-32301", "Insects, spiders, micro-organisms": "32301-32501", "Prehistoric animals": "32501-32502"}
 #########################################################################
 # Function to predict outputs
 def predict(img, model_name):
     if (model_name == 'Geometric Figures & Solids'):
         geo_learn = import_model('geo_model.pkl')
-        pred = geo_learn.predict(img)[0]
-        labels = geo
-        pred = str(pred)
-        fix_pred = pred[4:-4]
-        return fix_pred + ", Design Codes: " + labels[fix_pred]
+        preds = geo_learn.predict(img)
     elif (model_name == 'Scenery, Natural Phenomena'):
         landscape_learn = import_model('landscape_model.pkl')
-        pred = landscape_learn.predict(img)[0]
-        labels = scen
-        pred = str(pred)
-        fix_pred = pred[4:-4]
-        return fix_pred + ", Design Codes: " + labels[fix_pred]
+        preds = landscape_learn.predict(img)
     elif (model_name == 'Human & Supernatural Beings'):
         human_learn = import_model('human_model.pkl')
-        pred = human_learn.predict(img)[0]
-        labels = char
-        pred = str(pred)
-        fix_pred = pred[4:-4]
-        return fix_pred + ", Design Codes: " + labels[fix_pred]
+        preds = human_learn.predict(img)
     elif (model_name == 'Colors & Characters'):
         colors_learn = import_model('colors_model.pkl')
-        pred = colors_learn.predict(img)[0]
-        labels = misc
-        pred = str(pred)
-        fix_pred = pred[4:-4]
-        return fix_pred
+        preds = colors_learn.predict(img) 
     elif (model_name == 'Buildings, Dwellings & Furniture'):
         build_learn = import_model('buildings_model.pkl')
-        pred = build_learn.predict(img)[0]
-        labels = build
-        pred = str(pred)
-        fix_pred = pred[4:-4]
-        return fix_pred
+        preds = build_learn.predict(img)
     elif (model_name == 'Animals'):
-        anim_learn = import_model('animals_model.pkl')
-        pred = anim_learn.predict(img)[0]
-        labels = anim
-        pred = str(pred)
-        fix_pred = pred[4:-4]
-        return fix_pred
+        anim_learn = import_model('animals.pkl')
+        preds = anim_learn.predict(img)
+
+    label_pred = str(preds[0])
+    acc = [float(preds[2][i]) for i in range(len(preds[2])) if preds[2][i] > 0.5]
+    pred = {label_pred: float(acc[0])}
+        
+    return pred
 #########################################################################
 title = "TM-TKO Trademark Logo Image Classification Model"
 description = "Users can upload an image and corresponding image file name to get US design-code standard predictions on a trained model that utilizes the benchmark ResNet50 architecture."