From b8b4288d29d68a801e4b3a07f529c7704914bdfc Mon Sep 17 00:00:00 2001
From: elenaverbitskaya <44786267+elenaverbitskaya@users.noreply.github.com>
Date: Fri, 8 Jan 2021 20:06:40 +0200
Subject: [PATCH] Update generate_tfrecord.py

---
 generate_tfrecord.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/generate_tfrecord.py b/generate_tfrecord.py
index e128b14..93346e8 100755
--- a/generate_tfrecord.py
+++ b/generate_tfrecord.py
@@ -28,8 +28,24 @@
 
 # TO-DO replace this with label map
 def class_text_to_int(row_label):
-    if row_label == 'apple':
+    if row_label == 'cocacola_violet_balls':
         return 1
+    elif row_label == 'cocacola_yellow_dragons':
+        return 2
+    elif row_label == 'cocacola_blue_flowers':
+        return 3
+    elif row_label == 'cocacola_green_apples':
+        return 4
+    elif row_label == 'cocacola_yellow_pattern':
+        return 5
+    elif row_label == 'cocacola_white_lines':
+        return 6
+    elif row_label == 'cocacola_blue_fish':
+        return 7
+    elif row_label == 'cocacola_pink_feather':
+        return 8
+    elif row_label == 'cocacola_pink_violet_feather':
+        return 9
     else:
         return 0