implementing app

Object_Smiles.py

@@ -0,0 +1,157 @@
+from torch.utils.data import Dataset, DataLoader, Subset
+from robust_detection.data_utils.rcnn_data_utils import *
+import pytorch_lightning as pl
+import robust_detection.transforms as T
+
+DATA_FOLDER = os.path.join(os.path.dirname(__file__))
+def get_transform():
+        transforms = []
+        transforms.append(T.ToTensor())
+        return T.Compose(transforms)
+
+class Objects_Smiles(pl.LightningDataModule):
+    def __init__(self, data_path, **kwargs):
+                super().__init__()
+                self.batch_size = 1
+                self.num_workers = 4
+                self.data_path = data_path
+                self.transforms = get_transform()
+                self.base_class = Objects_Detection_Predictor_Dataset
+    def prepare_data(self):
+                dataset = self.base_class(os.path.join(DATA_FOLDER, self.data_path), self.transforms)
+                self.train = dataset
+                self.test  = dataset
+                self.val   = dataset
+                
+                self.test_ood = dataset
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.train,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            drop_last=False,
+            pin_memory=True,
+            collate_fn=collate_tuple
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.val,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            drop_last=False,
+            pin_memory=True,
+            collate_fn=collate_tuple
+        )
+
+    def test_dataloader(self):
+        return DataLoader(
+            self.test,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            drop_last=False,
+            pin_memory=True,
+            collate_fn=collate_tuple
+        )
+
+    def test_ood_dataloader(self, shuffle=False):
+        return DataLoader(
+            self.test_ood,
+            batch_size=self.batch_size,
+            shuffle=shuffle,
+            num_workers=self.num_workers,
+            drop_last=False,
+            pin_memory=True,
+            collate_fn=collate_tuple
+        )
+   
+    @classmethod
+    def add_dataset_specific_args(cls, parent):
+        import argparse
+        parser = argparse.ArgumentParser(parents=[parent], add_help=False)
+        parser.add_argument('--data_path', type=str,
+                            default="mnist/alldigits/")
+        return parser
+
+
+class Objects_fold_Smiles(pl.LightningDataModule):
+    def __init__(self, data_path, fold, **kwargs):
+                super().__init__()
+                self.batch_size = 1
+                self.num_workers = 4
+                self.data_path = data_path
+                self.fold = fold
+                self.transforms = get_transform()
+               # self.base_class = Objects_Detection_Predictor_Dataset
+                self.base_class = Objects_Detection_Dataset
+    def prepare_data(self):
+                dataset = self.base_class(os.path.join(DATA_FOLDER, self.data_path), self.transforms)
+                if self.fold > -1:
+                   train_idx = np.load(os.path.join(DATA_FOLDER, f"{self.data_path}", "../folds", str(self.fold), "train_idx.npy"))
+                   self.train = Subset(dataset, train_idx)
+                   val_idx = np.load(os.path.join(DATA_FOLDER, f"{self.data_path}", "../folds", str(self.fold), "val_idx.npy"))
+
+                   self.val = Subset(dataset, val_idx)
+                else:
+                   self.train = dataset
+                   self.val   = dataset
+                self.test = self.val
+                self.test_ood = self.test
+
+    def train_dataloader(self):
+        return DataLoader(
+            self.train,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            drop_last=False,
+            pin_memory=True,
+            collate_fn=collate_tuple
+        )
+
+    def val_dataloader(self):
+        return DataLoader(
+            self.val,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            drop_last=False,
+            pin_memory=True,
+            collate_fn=collate_tuple
+        )
+
+    def test_dataloader(self):
+        return DataLoader(
+            self.test,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            drop_last=False,
+            pin_memory=True,
+            collate_fn=collate_tuple
+        )
+    
+    def test_ood_dataloader(self, shuffle=False):
+        return DataLoader(
+            self.test_ood,
+            batch_size=self.batch_size,
+            shuffle=shuffle,
+            num_workers=self.num_workers,
+            drop_last=False,
+            pin_memory=True,
+            collate_fn=collate_tuple
+        )
+
+    @classmethod
+    def add_dataset_specific_args(cls, parent):
+        import argparse
+        parser = argparse.ArgumentParser(parents=[parent], add_help=False)
+        parser.add_argument('--data_path', type=str,
+                            default="mnist/alldigits/")
+        parser.add_argument('--fold', type=int,
+                            default=0)
+        return parser

app.py

@@ -7,7 +7,7 @@ import numpy as np
 #import pathlib
 #from AtomLenz import *
 #from utils_graph import *
-#from Object_Smiles import Objects_Smiles 
+from Object_Smiles import Objects_Smiles 
 
 #from robust_detection import wandb_config
 from robust_detection import utils
@@ -33,8 +33,9 @@ dir_list.sort(key=os.path.getctime, reverse=True)
 checkpoint_file_atoms = [f for f in dir_list if "ckpt" in f][0]
 model_atom = model_cls.load_from_checkpoint(checkpoint_file_atoms)
 model_atom.model.roi_heads.score_thresh = 0.65 
-
-
+data_cls = Objects_Smiles
+dataset = data_cls(datapath="./uploads/")
+#    dataset.prepare_data()
 st.title("Atom Level Entity Detector")
 
 image_file = st.file_uploader("Upload a chemical structure candidate image",type=['png','jpeg','jpg'])
@@ -45,7 +46,10 @@ if image_file is not None:
    image = Image.open(image_file)
    col1.image(image, use_column_width=True)
    with open(os.path.join("uploads",image_file.name),"wb") as f: 
-        f.write(image_file.getbuffer())         
+        f.write(image_file.getbuffer()) 
    st.success("Saved File")
+   dataset.prepare_data()
+   trainer = pl.Trainer(logger=False)
+   atom_preds = trainer.predict(model_atom, dataset.test_dataloader())
 x = st.slider('Select a value')
 st.write(x, 'squared is', x * x)