Testing whole program it 1

app.py

@@ -4,13 +4,86 @@ import streamlit as st
 from datasets import load_dataset
 from PIL import Image
 from matplotlib.pyplot import imshow
+from transformers import pipeline
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
 
 def load_dataset_from_Hugging_Face():
   reloaded_dataset_thermal = load_dataset("gokulraja17/rice-thermal-demo")
   reloaded_dataset_rgb = load_dataset("gokulraja17/rice-rgb-demo2")
   return reloaded_dataset_thermal,reloaded_dataset_rgb
 
+def load_model_from_Hugging_Face():
+  feature_extractor_thermal = AutoFeatureExtractor.from_pretrained("gokulraja17/convnext-fine-tuned-thermal-demo2")
+  model_thermal = AutoModelForImageClassification.from_pretrained("gokulraja17/convnext-fine-tuned-thermal-demo2")
+  feature_extractor_rgb = AutoFeatureExtractor.from_pretrained("gokulraja17/convnext-fine-tuned-rgb-demo2")
+  model_rgb = AutoModelForImageClassification.from_pretrained("gokulraja17/convnext-fine-tuned-rgb-demo2")
+  return feature_extractor_thermal,model_thermal,feature_extractor_rgb,model_rgb
+
+def get_image_and_label(reloaded_dataset_thermal,reloaded_dataset_rgb,index):
+  url_thermal = reloaded_dataset_thermal["test"][index]
+  url_rgb = reloaded_dataset_rgb["test"][index]
+  return url_thermal["image"],url_thermal["label"],url_rgb["image"],url_rgb["label"]
+
+def create_pipeline(feature_extractor_thermal,model_thermal,feature_extractor_rgb,model_rgb):
+
+  pipe_rgb = pipeline("image-classification",
+                model=model_rgb,
+                feature_extractor=feature_extractor_rgb)
+  pipe_thermal = pipeline("image-classification",
+                model=model_thermal,
+                feature_extractor=feature_extractor_thermal)
+  return pipe_thermal,pipe_rgb
+
+def get_output_label(result):
+  max = -999
+  label = ""
+  for i in result:
+    if(max<i['score']):
+      max = i['score']
+      label = i["label"]
+  return label
+
 dataset_thermal,dataset_rgb = load_dataset_from_Hugging_Face()
+labels=dataset_thermal["train"].features["label"].names
+id2label={k:v for k,v in enumerate(labels)}
+print(id2label)
+label2id = {v:k for k,v in enumerate(labels)}
+print(label2id)
+
+feature_extractor_thermal,model_thermal,feature_extractor_rgb,model_rgb = load_model_from_Hugging_Face()
+index=3
+image_thermal,thermal_label,image_rgb,rgb_label = get_image_and_label(dataset_thermal,dataset_rgb,index)
+print("Thermal Input...")
+imshow(np.asarray(image_thermal))
+print(id2label[thermal_label])
+print("RGB Input...")
+imshow(np.asarray(image_rgb))
+print(id2label[rgb_label])
+
+def consolidate_score(thermal_result,rgb_result):
+  consol_score_1=[]
+  consol_score_2=[]
+  consol_score=[]
+  for i in range(0,len(thermal_result)):
+    consol_score_1.append(list(thermal_result[i].values()))
+    consol_score_2.append(list(rgb_result[i].values()))
+  temp_label=""
+  for i in range(0,len(thermal_result)):
+    temp_label=consol_score_1[i][1]
+    for j in range(0,len(rgb_result)):
+      if(temp_label==consol_score_2[j][1]):
+        consol_score.append([temp_label,consol_score_1[i][0],consol_score_2[j][0]])
+  output_consol=[]
+  for i in consol_score:
+    output_consol.append({'score':(((i[2]/100)*35)+((i[1]/100)*65)),'label':i[0]})
+  return output_consol
+
+pipe_thermal,pipe_rgb = create_pipeline(feature_extractor_thermal,model_thermal,feature_extractor_rgb,model_rgb)
+thermal_result=pipe_thermal(image_thermal)
+rgb_result=pipe_rgb(image_rgb)
+consolidated_result=get_output_label(consolidate_score(thermal_result,rgb_result))
+print(consolidated_result)
+
 uploaded_file = st.file_uploader("Choose a image file", type=["jpg","png"])
 
 if uploaded_file is not None: