add new gallery view, add selectable images

app.py

@@ -18,7 +18,7 @@ class GalleryApp:
         self.promptBook = promptBook
         st.set_page_config(layout="wide")
 
-    def gallery(self, items, col_num, info):
+    def gallery_masonry(self, items, col_num, info):
         cols = st.columns(col_num)
         # # sort items by brisque score
         # items = items.sort_values(by=['brisque'], ascending=True).reset_index(drop=True)
@@ -34,23 +34,49 @@ class GalleryApp:
                 #         st.image(image, use_column_width=True)
                 #     with tab2:
                 #         st.image(image, use_column_width=True)
+
+                # show checkbox
+                self.promptBook.loc[items.iloc[idx]['row_idx'].item(), 'checked'] = st.checkbox(
+                    'Select', value=self.promptBook.loc[items.iloc[idx]['row_idx'].item(), 'checked'],
+                    key=f'select_{idx}')
+
                 for key in info:
                     st.write(f"**{key}**: {items.iloc[idx][key]}")
 
+    def gallery_standard(self, items, col_num, info):
+        rows = len(items) // col_num + 1
+        containers = [st.container() for _ in range(rows*2)]
+        for idx in range(0, len(items), col_num):
+            # assign one container for each row
+            row_idx = (idx // col_num) * 2
+            with containers[row_idx]:
+                cols = st.columns(col_num)
+                for j in range(col_num):
+                    if idx + j < len(items):
+                        with cols[j]:
+                            # show image
+                            image = st.session_state.images[items.iloc[idx+j]['row_idx'].item()]['image']
+                            st.image(image,
+                                     use_column_width=True,
+                            )
+
+                            # show checkbox
+                            self.promptBook.loc[items.iloc[idx+j]['row_idx'].item(), 'checked'] = st.checkbox('Select', value=self.promptBook.loc[items.iloc[idx+j]['row_idx'].item(), 'checked'] , key=f'select_{idx+j}')
+
+                            # show selected info
+                            for key in info:
+                                st.write(f"**{key}**: {items.iloc[idx+j][key]}")
+
+                            # st.write(row_idx/2, idx+j, rows)
+                            # extra_info = st.checkbox('Extra Info', key=f'extra_info_{idx+j}')
+                            # if extra_info:
+                            #     with containers[row_idx+1]:
+                            #         st.image(image, use_column_width=True)
+
     def app(self):
         st.title('Model Coffer Gallery')
         st.write('This is a gallery of images generated by the models in the Model Coffer')
 
-        # metadata, images = st.columns([1, 3])
-        # with images:
-            # prompt_tags = self.promptBook['tag'].unique()
-            # # sort tags by alphabetical order
-            # prompt_tags = np.sort(prompt_tags)
-            #
-            # selecters = st.columns(3)
-            # with selecters[0]:
-            #     tag = st.selectbox('Select a tag', prompt_tags)
-
         with st.sidebar:
             prompt_tags = self.promptBook['tag'].unique()
             # sort tags by alphabetical order
@@ -96,11 +122,11 @@ class GalleryApp:
 
                 res = requests.get(f'https://civitai.com/images', params={'post_id': prompt_id})
                 st.write(res)
-                # image_url = res.json()['items'][0]['url']
-                # st.image(image_url, use_column_width=True)
+                image_url = res.json()['items'][0]['url']
+                st.image(image_url, use_column_width=True)
 
         # with images:
-        selecters = st.columns([1, 1, 2])
+        selecters = st.columns([1, 1, 2, 0.5])
 
         with selecters[0]:
             # sort_by = st.selectbox('Sort by', items.columns[11: -1])
@@ -116,6 +142,7 @@ class GalleryApp:
                 order = False
 
         items = items.sort_values(by=[sort_by], ascending=order).reset_index(drop=True)
+
         with selecters[2]:
             info = st.multiselect('Show Info',
                                   ['model_download_count', 'model_name', 'model_id',
@@ -123,7 +150,46 @@ class GalleryApp:
                                   default=sort_by)
 
         col_num = st.slider('Number of columns', min_value=1, max_value=9, value=4, step=1, key='col_num')
-        self.gallery(items, col_num, info)
+
+        with selecters[3]:
+            filter = st.selectbox('Filter', ['All', 'Checked', 'Unchecked'])
+            if filter == 'Checked':
+                items = items[items['checked'] == True].reset_index(drop=True)
+            elif filter == 'Unchecked':
+                items = items[items['checked'] == False].reset_index(drop=True)
+
+        with st.form(key='my_form'):
+            buttons = st.columns([1, 1, 1])
+            with buttons[0]:
+                submit = st.form_submit_button('Save selections', on_click=self.save_checked, use_container_width=True, type='primary')
+            with buttons[1]:
+                submit = st.form_submit_button('Reset current prompt', on_click=self.reset_current_prompt, kwargs={'prompt_id': prompt_id} , use_container_width=True)
+            with buttons[2]:
+                submit = st.form_submit_button('Reset all selections', on_click=self.reset_all, use_container_width=True)
+            self.gallery_standard(items, col_num, info)
+
+    def reset_current_prompt(self, prompt_id):
+        # reset current prompt
+        self.promptBook.loc[self.promptBook['prompt_id'] == prompt_id, 'checked'] = False
+        self.save_checked()
+
+    def reset_all(self):
+        # reset all
+        self.promptBook.loc[:, 'checked'] = False
+        self.save_checked()
+
+    def save_checked(self):
+        # save checked images to huggingface dataset
+        dataset = load_dataset('NYUSHPRP/ModelCofferMetadata', split='train')
+        # get checked images
+        checked_info = self.promptBook['checked']
+
+        if 'checked' in dataset.column_names:
+            dataset = dataset.remove_columns('checked')
+        dataset = dataset.add_column('checked', checked_info)
+
+        # print('metadata dataset: ', dataset)
+        dataset.push_to_hub('NYUSHPRP/ModelCofferMetadata', split='train')
 
 
 if __name__ == '__main__':
@@ -140,6 +206,10 @@ if __name__ == '__main__':
         print('loading promptBook')
 
         st.session_state.promptBook = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferMetadata', split='train'))
+        # add 'checked' column to promptBook if not exist
+        if 'checked' not in st.session_state.promptBook.columns:
+            st.session_state.promptBook['checked'] = False
+
         st.session_state.images = load_from_disk(os.path.join(os.getcwd(), 'data', 'promptbook'))
         print(st.session_state.images)
         print('images loaded')