Update

app.py

@@ -42,20 +42,16 @@ def solve(file_obj):
                 total += X * cycle
             X += int(line.split()[-1])
             sprite = X
-        if part_2:
-            # Go back to first position after 40 pixels
-            cycle = cycle % DISPLAY_WIDTH
+        # Go back to first position after 40 pixels
+        cycle = cycle % DISPLAY_WIDTH
 
-    print("Solution to part 1: ", total)
+    # print("Solution to part 1: ", total)
     res = "\n".join(["".join(crt[i:i+DISPLAY_WIDTH]) for i in range(0, len(crt), DISPLAY_WIDTH)])
-    print("Solution to part 2:")
     print(res)
-
-
+    
     # Bonus: make an image and then use OCR
     # to extract the text.
-    img = np.array(crt)
-
+    img = np.array([[c.replace(PIXEL, 1).replace(EMPTY, 0) for c in crt[i:i+DISPLAY_WIDTH]] for i in range(0, len(crt), DISPLAY_WIDTH)) ])
     plt.imshow(img, cmap="binary")
     plt.axis('off')
     plt.savefig("day_10.png")